import math
import numpy as np
import matplotlib.pyplot as plt

# --- 1. 定义常量和坐标 ---
SQRT3 = math.sqrt(3)
A2 = np.array([250000, 500000 + 100000 * SQRT3])
A3 = np.array([300000, 500000 + 50000 * SQRT3])
A4 = np.array([300000, 50000 * SQRT3])
A5 = np.array([250000, 0])
AOI = np.array([A2, A3, A4, A5])
A7 = np.array([0, 50000 * SQRT3])
A8 = np.array([0, 500000 + 50000 * SQRT3])

# 关键约束和参数
NUM_UAVS = 5
UAV_MAX_SPEED = 150
MAX_FLIGHT_TIME_S = 7200

# --- 2. 辅助函数 ---
def get_y_boundaries_for_centerline(x):
    y_bottom = SQRT3 * (x - A5[0])
    y_top = A2[1] - SQRT3 * (x - A2[0])
    return y_bottom, y_top

# --- 3. 主规划流程 ---
def plan_final_mission_corrected():
    print("--- Final Corrected Mission Plan (Clustered Initial Positions) ---")
    
    # 1. 定义唯一的中心扫描线
    centerline_x = A5[0] + (A3[0] - A5[0]) / 2  # X = 275000
    y_start_total, y_end_total = get_y_boundaries_for_centerline(centerline_x)
    total_scan_length = y_end_total - y_start_total
    
    # 2. 将扫描线切分为5段
    segment_length = total_scan_length / NUM_UAVS

    # 3. *** 这是核心修正：正确计算无人机初始位置 ***
    print("Calculating correct clustered initial positions around the midpoint...")
    uav_initial_positions = []
    midpoint_A7A8 = (A7 + A8) / 2
    spacing = 5000  # 5km

    for i in range(NUM_UAVS):
        # 计算每个无人机相对于中点的偏移量
        # i=0 -> -2 * spacing; i=1 -> -1 * spacing; i=2 -> 0; i=3 -> 1 * spacing; i=4 -> 2 * spacing
        offset_y = (i - (NUM_UAVS - 1) / 2) * spacing
        pos = np.array([midpoint_A7A8[0], midpoint_A7A8[1] + offset_y])
        uav_initial_positions.append(pos)
    
    # 为了任务分配的逻辑性，我们将初始点从上到下排序
    uav_initial_positions.sort(key=lambda p: p[1], reverse=True)
    print("Initial positions corrected and sorted.")

    # 4. 为每架无人机规划轨迹
    all_trajectories = []
    for i in range(NUM_UAVS):
        # 从上到下分配线段给从上到下的无人机
        segment_y_start = y_end_total - (i + 1) * segment_length
        segment_y_end = y_end_total - i * segment_length
        
        start_wp = uav_initial_positions[i]
        ingress_wp = np.array([centerline_x, segment_y_end])
        egress_wp = np.array([centerline_x, segment_y_start])
        
        trajectory = np.array([start_wp, ingress_wp, egress_wp, start_wp])
        all_trajectories.append(trajectory)
        
    return all_trajectories, uav_initial_positions

# --- 4. 分析与可视化 ---
def analyze_trajectories(trajectories, speed_ms):
    analysis_results = []
    print("\n--- Final Mission Analysis (Corrected) ---")
    print("-" * 70)
    print(f"{'UAV ID':<10} {'Total Distance (km)':<25} {'Est. Time (minutes)':<25} {'Status':<10}")
    print("-" * 70)
    for i, traj in enumerate(trajectories):
        distance = np.sum(np.linalg.norm(np.diff(traj, axis=0), axis=1))
        time_min = (distance / speed_ms) / 60
        status = "OK" if time_min <= MAX_FLIGHT_TIME_S / 60 else "FAIL"
        print(f"UAV {i+1:<7} {distance/1000:<25.2f} {time_min:<25.2f} {status:<10}")
        analysis_results.append({"status": status})
    print("-" * 70)
    return analysis_results

# --- 主执行流程 ---
final_trajectories, final_initial_pos = plan_final_mission_corrected()
final_analysis = analyze_trajectories(final_trajectories, UAV_MAX_SPEED)

if any(res['status'] == 'FAIL' for res in final_analysis):
    print("\nWARNING: Mission as planned is not feasible.")
else:
    print("\nSUCCESS: All UAV trajectories are feasible within the 120-minute limit.")

# 可视化
plt.figure(figsize=(12, 16))
# 绘制背景和目标区域
A7A8_line = np.array([A7, A8])
plt.plot(A7A8_line[:,0], A7A8_line[:,1], 'k:', label='A7-A8 Line')
AOI_poly = np.array([A2, A3, A4, A5, A2])
plt.fill(AOI_poly[:, 0], AOI_poly[:, 1], 'lightgray', alpha=0.6, label='Target Area')
# 绘制【正确】的初始点
initial_np = np.array(final_initial_pos)
plt.scatter(initial_np[:, 0], initial_np[:, 1], c='red', marker='s', s=100, zorder=5, label='UAV Clustered Start Positions')

# 绘制轨迹
colors = plt.cm.viridis(np.linspace(0, 1, NUM_UAVS))
for i, traj in enumerate(final_trajectories):
    plt.plot(traj[:, 0], traj[:, 1], color=colors[i], label=f'UAV {i+1} Trajectory', marker='o')
    plt.text(traj[1, 0] + 5000, (traj[1, 1] + traj[2, 1]) / 2, f'UAV {i+1}', color=colors[i], weight='bold')

# 清理图例
handles, labels = plt.gca().get_legend_handles_labels()
by_label = dict(zip(labels, handles))
plt.legend(by_label.values(), by_label.keys())

plt.title('Final Corrected Plan: 5-UAV Segmented Line Scan', fontsize=16)
plt.xlabel('X Coordinate (meters)')
plt.ylabel('Y Coordinate (meters)')
plt.grid(True)
plt.axis('equal')
plt.show()